/*
 * Copyright (c) 2005, 2014 vacoor
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 */
package org.vacoor.nothing.security.shiro.filter;

import org.apache.shiro.authc.AuthenticationException;
import org.apache.shiro.authc.AuthenticationToken;
import org.apache.shiro.codec.Base64;
import org.apache.shiro.session.Session;
import org.apache.shiro.subject.Subject;
import org.apache.shiro.web.filter.AccessControlFilter;
import org.apache.shiro.web.filter.authc.FormAuthenticationFilter;
import org.apache.shiro.web.servlet.Cookie;
import org.apache.shiro.web.servlet.SimpleCookie;
import org.apache.shiro.web.util.SavedRequest;
import org.apache.shiro.web.util.WebUtils;
import org.vacoor.nothing.common.util.Bytes;

import javax.crypto.Cipher;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * redirect=1&redirectUrl=&style=mini
 * https://login.taobao.com/member/login.jhtml?style=mini&redirectURL=
 * @author vacoor
 */
public class AuthenticationFilter extends FormAuthenticationFilter {
    public static final String PRIVATE_KEY = AuthenticationFilter.class + ".PRIVATE_KEY";
    private static final int BIT_SIZE = 1024;
    public static final boolean DEFAULT_ENABLED_RSA_PASSWORD = true;
    public static final String DEFAULT_PUBLIC_KEY_COOKIE_NAME = "npk";
    public static final String DEFAULT_REDIRECT_PARAM = "redirect";

    private String redirectParam = DEFAULT_REDIRECT_PARAM;
    private boolean enabledRsaPassword = DEFAULT_ENABLED_RSA_PASSWORD;
    private Cookie cookie;

    public AuthenticationFilter() {
        Cookie cookie = new SimpleCookie(DEFAULT_PUBLIC_KEY_COOKIE_NAME);
        cookie.setHttpOnly(false);

        this.cookie = cookie;
    }

    @Override
    protected boolean onAccessDenied(ServletRequest request, ServletResponse response) throws Exception {
        boolean allowed = super.onAccessDenied(request, response);

        if (enabledRsaPassword && isLoginRequest(request, response)) {
            // if request login page or (login submit but login failed)
            if (!isLoginSubmission(request, response) || (isLoginSubmission(request, response) && !allowed)) {
                KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
                generator.initialize(BIT_SIZE);
                KeyPair keyPair = generator.generateKeyPair();
                RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
                Session session = getSubject(request, response).getSession();
                session.setAttribute(PRIVATE_KEY, privateKey);

                RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
                BigInteger modulus = publicKey.getModulus();
                BigInteger exponent = publicKey.getPublicExponent();

                String base64 = Base64.encodeToString(modulus.toByteArray()) + ":" + Base64.encodeToString(exponent.toByteArray());
                Cookie cookie = new SimpleCookie(this.cookie);
                cookie.setValue(base64);
                cookie.saveTo((HttpServletRequest) request, (HttpServletResponse) response);
            }
        }

        return allowed;
    }

    @Override
    protected boolean onLoginSuccess(AuthenticationToken token, Subject subject, ServletRequest request, ServletResponse response) throws Exception {
        // reset session when login success
        resetSession(subject);
        cookie.removeFrom((HttpServletRequest) request, (HttpServletResponse) response);
        return super.onLoginSuccess(token, subject, request, response);
    }

    @Override
    protected void issueSuccessRedirect(ServletRequest request, ServletResponse response) throws Exception {
        // super.issueSuccessRedirect(request, response);

        String successUrl = null;
        boolean contextRelative = true;
        SavedRequest savedRequest = WebUtils.getAndClearSavedRequest(request);
        if (savedRequest != null && savedRequest.getMethod().equalsIgnoreCase(AccessControlFilter.GET_METHOD)) {
            successUrl = savedRequest.getRequestUrl();
            contextRelative = false;
        }

        if (successUrl == null) {
            successUrl = getSuccessUrl();
        }

        if (successUrl == null) {
            throw new IllegalStateException("Success URL not available via saved request or via the " +
                    "successUrlFallback method parameter. One of these must be non-null for " +
                    "issueSuccessRedirect() to work.");
        }

        if (contextRelative) {
            successUrl = ((HttpServletRequest) request).getContextPath() + successUrl;
        }

        boolean redirect = getRedirect(request);
        if (!redirect) {
            response.getWriter().println("{\"href\":" + "\"" + successUrl + "\"}");
        } else {
            WebUtils.issueRedirect(request, response, successUrl);
        }
    }

    @Override
    protected AuthenticationToken createToken(ServletRequest request, ServletResponse response) {
        String username = getUsername(request);
        String password = getPassword(request);

        if (enabledRsaPassword) {
            try {
                RSAPrivateKey privateKey = (RSAPrivateKey) getSubject(request, response).getSession().getAttribute(PRIVATE_KEY);
                byte[] encoded = Base64.decode(password);
                Cipher cipher = Cipher.getInstance("RSA");
                cipher.init(Cipher.DECRYPT_MODE, privateKey);
                byte[] bytes = cipher.doFinal(encoded);
                password = Bytes.toString(bytes);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return createToken(username, password, request, response);
    }

    /**
     * 覆盖父类默认错误属性设置, 将异常message 作为值写入, 默认为异常完全限定名
     */
    @Override
    protected void setFailureAttribute(ServletRequest request, AuthenticationException ae) {
        request.setAttribute(getFailureKeyAttribute(), ae.getLocalizedMessage());
    }

    /**
     * 重置 subject session
     *
     * @param subject
     */
    protected void resetSession(Subject subject) {
        Session session = subject.getSession(false);
        if (session == null) {
            return;
        }

        Map<Object, Object> transfer = new HashMap<Object, Object>();
        Collection<Object> attrKeys = session.getAttributeKeys();
        for (Object key : attrKeys) {
            transfer.put(key, session.getAttribute(key));
        }

        /*
         * 重置 session 后会引起 URL 重写 ....;JSESSIONID=newSessionId
         * 在 spring mvc 中存在问题需要使用
         * @RequestMapping(value = {"/{successUrl:successUrl;?.*}"})
         * spring3.2.2 bug see  http://jinnianshilongnian.iteye.com/blog/1831408
         */
        session.stop();

        session = subject.getSession();
        for (Map.Entry<Object, Object> entry : transfer.entrySet()) {
            session.setAttribute(entry.getKey(), entry.getValue());
        }
    }

    protected boolean getRedirect(ServletRequest request) {
        return WebUtils.isTrue(request, getRedirectParam());
    }

    public String getRedirectParam() {
        return redirectParam;
    }

    public void setRedirectParam(String redirectParam) {
        this.redirectParam = redirectParam;
    }
}
